#include <iostream>
#include <cmath>
#include "alignment.h"
#include "MbMatrix.h"
#include "model.h"
#include "settings.h"
#include "tree.h"
#include "util.h"

using namespace std;



Model::Model(Alignment *ap, MbRandom *rp, Settings *sp) {

	alignmentPtr = ap;
	ranPtr       = rp;
	settingsPtr  = sp;
	
	/* set up the conditional likelihoods */
	int numTaxa      = alignmentPtr->getNumTaxa();
	    numNodes     = 2 * numTaxa - 2;
	int numChar      = alignmentPtr->getNumChar();
	int numStates    = 4;
	int condLikeSize = numNodes * numChar * numStates;
	cls = new double[condLikeSize];
	if ( !cls )
		{
		cerr << "ERROR: Problem allocating conditional likelihoods" << endl;
		exit(1);
		}
	for (int i=0; i<condLikeSize; i++)
		cls[i] = 0.0;
		
	/* allocate a matrix of pointers so we can quickly position ourselves in
	   the vector of conditional likelihoods (cls) */
	clsPtr = new double**[numNodes];
	clsPtr[0] = new double*[numNodes * numChar];
	for (int i=1; i<numNodes; i++)
		clsPtr[i] = clsPtr[i-1] + numChar;
	for (int i=0; i<numNodes; i++)
		for (int j=0; j<numChar; j++)
			clsPtr[i][j] = &cls[i * numChar * numStates + j * numStates];
		
	/* initialize the conditional likelihoods at the tips of the tree, which
	   are indexed 0, 1, ..., numTaxa-1 */
	for (int i=0; i<numTaxa; i++)
		{
		for (int j=0; j<numChar; j++)
			{
			double *clp = clsPtr[i][j];
			int nucCode = alignmentPtr->getNucleotide(i, j);
			int nucs[4];
			alignmentPtr->getPossibleNucs(nucCode, nucs);
			for (int k=0; k<numStates; k++)
				{
				if ( nucs[k] == 1 )
					clp[k] = 1.0;
				}
			}
		}
	
	/* make a tree */
	string treeStr = getLineFromFile( settingsPtr->getTreeFileName(), 0 ); // read the tree from the file
	tree = new Tree(alignmentPtr, ranPtr, treeStr);
	tree->print();
	
	/* set up transition probability matrices according to the Jukes-Cantor (1969) model of
	   DNA substitution */
	P = new MbMatrix<double>[numNodes];
	for (int i=0; i<numNodes; i++)
		P[i] = MbMatrix<double>(4, 4);
	for (int n=0; n<numNodes; n++)
		{
		Node *p = tree->getDownPassNode( n );
		if (p->getAnc() != NULL)
			{
			int idx = p->getIndex();
			double v = p->getV();
			double x = exp(-(4.0/3.0)*v);
			for (int i=0; i<4; i++)
				{
				for (int j=0; j<4; j++)
					{
					if (i == j)
						P[idx][i][j] = 0.25 + 0.75 * x;
					else
						P[idx][i][j] = 0.25 - 0.25 * x;
					}
				}
			}
		}
	
}

Model::~Model(void) {

	delete [] cls;
	delete [] clsPtr[0];
	delete [] clsPtr;
	delete tree;
	
}

double Model::calculateLikelihood(void) {

		/* Calculate conditional likelihoods down the tree. Note that
		   we have already calculated the transition probabilities in
		   the constructor. */
		for (int n=0; n<numNodes; n++)
			{
			Node *p = tree->getDownPassNode( n );
			if (p->getLft() != NULL && p->getRht() != NULL && p->getAnc() != NULL)
				{
				/* deal with a two-way split */
				if (p->getAnc()->getAnc() != NULL)
					{
					int lftIdx = p->getLft()->getIndex();
					int rhtIdx = p->getRht()->getIndex();
					int idx    = p->getIndex();
					double *clL = clsPtr[lftIdx][0];
					double *clR = clsPtr[rhtIdx][0];
					double *clP = clsPtr[idx   ][0];
					for (int c=0; c<alignmentPtr->getNumChar(); c++)
						{
						for (int i=0; i<4; i++)
							{
							double sumL = 0.0, sumR = 0.0;
							for (int j=0; j<4; j++)
								{
								sumL += clL[j] * P[lftIdx][i][j];
								sumR += clR[j] * P[rhtIdx][i][j];
								}
							clP[i] = sumL * sumR;
							}
						clL += 4;
						clR += 4;
						clP += 4;
						}
					}
				else
					{
					/* deal with the node to the left of the root node */
					int lftIdx = p->getLft()->getIndex();
					int rhtIdx = p->getRht()->getIndex();
					int ancIdx = p->getAnc()->getIndex();
					int idx    = p->getIndex();
					double *clL = clsPtr[lftIdx][0];
					double *clR = clsPtr[rhtIdx][0];
					double *clA = clsPtr[ancIdx][0];
					double *clP = clsPtr[idx   ][0];
					for (int c=0; c<alignmentPtr->getNumChar(); c++)
						{
						for (int i=0; i<4; i++)
							{
							double sumL = 0.0, sumR = 0.0, sumA = 0.0;
							for (int j=0; j<4; j++)
								{
								sumL += clL[j] * P[lftIdx][i][j];
								sumR += clR[j] * P[rhtIdx][i][j];
								sumA += clA[j] * P[idx   ][i][j];
								}
							clP[i] = sumL * sumR * sumA;
							}
						clL += 4;
						clR += 4;
						clA += 4;
						clP += 4;
						}
					}
				}
			}
			
	/* calculate likelihood */
	Node *p = tree->getRoot()->getLft();
	double *clP = clsPtr[p->getIndex()][0];
	double lnL = 0.0;
	for (int c=0; c<alignmentPtr->getNumChar(); c++)
		{
		double like = 0.0;
		for (int i=0; i<4; i++)
			like += clP[i] * 0.25;
		//cout << c << " -- " << log(like) << " " << alignmentPtr->getNumOfPattern(c) << endl;
		lnL += log(like) * alignmentPtr->getNumOfPattern(c);
		clP += 4;
		}
	
	return lnL;

}